// "Copyright [2021] <Copyright Shengkai Lin>"

// This file is some useful functions
#ifndef BASIC_H_
#define BASIC_H_

#include <omnetpp.h>
#include <vector>
#include <unordered_map>
#include <memory>
#include "network_controller/topology.h"

namespace ai_sim {

using namespace std;

typedef vector<pair<int, int>> IntPairVector;
typedef map<int, int> IntMap;

enum TaskType {
    RingAllreduce,
    DoubleBinaryTree,
    Butterfly
};

enum ClusterType {
    OpticalCluster,
    ElectricalCluster,
    IbCluster
};

enum SpecialRequireofFlow {
    NoRequirement,
    NeedControllerCooperate,  // Need to inform the controller before send
    NeedCompute  // Need to wait to compute
};

struct RouteInfo {
    int inPort;
    int outPort;
    int srcIdx;
    int destIdx;
};

typedef vector<ai_sim::RouteInfo> RouteVec;

struct Task {
    int taskId;
    TaskType tasktype;
    int numGpus;
    int numBatches;
    int modelsize;
    double compute_time;
};

struct Flow {
    int flowId;
    Task the_task;
    double data_size;
    int srcIdx;
    int destIdx;
    int outportId;
    int numUpstreamflows;
    vector<int> downstream_flows;
    map<int, RouteVec> dependRoute;
    // the bool represents that if the require is completed
    unordered_map<SpecialRequireofFlow, bool> requires;
    int offset = 0;
    bool isLastPacket = false;
    // vector<SpecialRequireofFlow> requires;
};
typedef std::shared_ptr<ai_sim::Flow> SharedFlowPtr;

typedef vector<Flow> FlowVector;

enum ProtocolType {
    TCP,
    IB,
    PFC
};

// Use in controller
typedef Topology::Node topoNode;
typedef Topology::Link topoLink;

typedef pair<topoNode*, topoNode*> comm_pair;

// Describe the 'stage': that is th commpair, flows, routes
// of a task stage. So Ring has only 1 stage, butterfly has
// multiple stages.
// NOTE : DONT CHANGE THE STRUCT CASUALLY!
struct stageboard {
    // Commpairs of different stages
    vector<comm_pair> comm_pairs;
    // Links of different stages
    vector<topoLink*> used_links;
    // The key of the maps below are the GPUIdx
    // Flows of different stages
    map<int, vector<Flow>> flows_on_gpu;
    // The key of the maps below are the Module Id
    // Routes of different stages
    map<int, vector<RouteInfo>> routes_on_switch;
};

// The different stage means that the routes or flows has modified
// So the ring only has 1 stage, and HD have many stages
struct taskboard {
    Task the_task;
    int flows_count;
    omnetpp::simtime_t start_time;
    int working_workers;
    vector<topoNode*> used_gpus;
    vector<stageboard> stage_info;
};

// Return the coordinate in the ROW/COLUMN
void convert2RL(int order, int side_length, int *r);

// Return the coordinate in the NODE/GPU
void convert2NG(int order, int *r);

// Return the nvlink port according to the src &dest
int compute_nvpeer(int src, int dest);

int get_rand(int biggest);

int compute_awgr_dest(int src, int dest, int side_length);

int get_dest_addr(const char *pkt_name);

int log2(int n);

// Get the output port to the specific
int getOutPortTo(omnetpp::cModule*, omnetpp::cModule*);

}  // namespace ai_sim

#endif  // BASIC_H_
